package main;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;

import translationModels.BidirectedModel;
import translationModels.BonusModel;
import translationModels.IBMModel;
import translationModels.IBMModel1;
import translationModels.IBMModel2;
import translationModels.TParam;
import translationModels.TParamsInitialization;

public class Main
{
	public static void main(String[] args)
	{
		InputArgs input = parseArgs(args);
		if (input == null)
			return;
		
		TrainingData trainingData;
		try
		{
			trainingData = new TrainingData(input.fFile, input.eFile, input.sentencesNo);
		}
		catch (IOException e)
		{
			System.out.println("There is an error in one of the input files");
			return;
		}
		
		TrainingData trainingDataReverse = null;
		if (input.bothDirections)
		{
			try
			{
				trainingDataReverse = new TrainingData(input.eFile, input.fFile, input.sentencesNo);
			}
			catch (IOException e)
			{
				System.out.println("There is an error in one of the input files");
				return;
			}
		}
		
		IBMModel model;
		if (input.modelNo == 1)
		{
			if (input.bothDirections)
				model = new BidirectedModel(new IBMModel1(trainingData), new IBMModel1(trainingDataReverse));
			else
				model = new IBMModel1(trainingData);
		}
		else if (input.modelNo == 2)
		{
			if (input.bothDirections)
				model = new BidirectedModel(new IBMModel2(trainingData), new IBMModel2(trainingDataReverse));
			else
				model = new IBMModel2(trainingData);
		}
		else
		{
			if (input.bothDirections)
				model = new BidirectedModel(new BonusModel(trainingData), new BonusModel(trainingDataReverse));
			else
				model = new BonusModel(trainingData);
		}
		
		HashMap<TParam, Double> tParams = null;
		if (input.randomInit)
			tParams = model.learnParams(TParamsInitialization.Random, input.iterations);
		else if (input.modelNo == 1)
			tParams = model.learnParams(TParamsInitialization.Uniform, input.iterations);
		else
			tParams = model.learnParams(TParamsInitialization.Bootstrapping, input.iterations);
		
		ArrayList<HashMap<Integer, Integer>> alignments = model.calculateAlignments();
		
		OutputData outputData = new OutputData(trainingData, tParams, alignments);
		try
		{
			outputData.writeOutput(input.outputFile);
			outputData.writeTParams(input.tFile);
		}
		catch (IOException e)
		{
			System.out.println("There is an error in one of the output files");
		}
	}
	
	private static InputArgs parseArgs(String[] args)
	{
		InputArgs input = new InputArgs();
		for (int i = 0; i < args.length; i++)
		{
			if (args[i].equals("-e"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the e file after the -e");
					return null;
				}
				
				File file = new File(args[i + 1]);
				if ((!file.exists()) || (!file.canRead()))
				{
					System.out.println("Cannot read from the e file");
					return null;
				}
				
				input.eFile = args[i + 1];
				i++;
			}
			
			else if (args[i].equals("-f"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the f file after the -f");
					return null;
				}
				
				File file = new File(args[i + 1]);
				if ((!file.exists()) || (!file.canRead()))
				{
					System.out.println("Cannot read from the f file");
					return null;
				}
				
				input.fFile = args[i + 1];
				i++;
			}
			
			else if (args[i].equals("-m"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the model number after the -m");
					return null;
				}
				
				int number;
				try
				{
					number = Integer.parseInt(args[i + 1]);
				}
				catch (NumberFormatException e)
				{
					System.out.println("The model number should be 1 or 2 or 3");
					return null;
				}
				
				if ((number != 1) && (number != 2) && (number != 3))
				{
					System.out.println("The model number should be 1 or 2 or 3");
					return null;
				}
				
				input.modelNo = number;
				i++;
			}
			
			else if (args[i].equals("-s"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the sentences number after the -s");
					return null;
				}
				
				int number;
				try
				{
					number = Integer.parseInt(args[i + 1]);
				}
				catch (NumberFormatException e)
				{
					System.out.println("The sentences number should be between 1 to 100,000");
					return null;
				}
				
				if ((number < 1) || (number > 100000))
				{
					System.out.println("The sentences number should be between 1 to 100,000");
					return null;
				}
				
				input.sentencesNo = number;
				i++;
			}
			
			else if (args[i].equals("-i"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the iterations number after the -i");
					return null;
				}
				
				int number;
				try
				{
					number = Integer.parseInt(args[i + 1]);
				}
				catch (NumberFormatException e)
				{
					System.out.println("The iterations number should be positive");
					return null;
				}
				
				if (number < 1)
				{
					System.out.println("The iterations number should be positive");
					return null;
				}
				
				input.iterations = number;
				i++;
			}
			
			else if (args[i].equals("-r"))
				input.randomInit = true;
			
			else if (args[i].equals("-b"))
				input.bothDirections = true;
			
			else if (args[i].equals("-o"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the output file after the -o");
					return null;
				}
				
				File file = new File(args[i + 1]);
				if ((file.exists()) && (!file.canWrite()))
				{
					System.out.println("Cannot write to the output file");
					return null;
				}
				
				input.outputFile = args[i + 1];
				i++;
			}
			
			else if (args[i].equals("-t"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the t params file after the -t");
					return null;
				}
				
				File file = new File(args[i + 1]);
				if ((file.exists()) && (!file.canWrite()))
				{
					System.out.println("Cannot write to the t params file");
					return null;
				}
				
				input.tFile = args[i + 1];
				i++;
			}
			
			else
			{
				System.out.println("Invalid argument: " + args[i]);
				return null;
			}
		}
		
		if ((input.eFile == null) || (input.fFile == null))
		{
			System.out.println("Wrong arguments");
			System.out.println("Use the following form:");
			System.out.println("-e e_file -f f_file [ -m model_no ] [ -s sentences_no ] [ -i iterations_no ] [ -r ] [ -b] [ -o output_file ] [ -t t_params_file ]");
			return null;
		}
		
		// Defaults
		if (input.modelNo == null)
			input.modelNo = 1;
		if (input.sentencesNo == null)
			input.sentencesNo = 100;
		if (input.iterations == null)
			input.iterations = 10;
		if (input.randomInit == null)
			input.randomInit = false;
		if (input.bothDirections == null)
			input.bothDirections = false;
		if (input.outputFile == null)
			input.outputFile = "output.txt";
		if (input.tFile == null)
			input.tFile = "tParams.txt";
		
		return input;
	}
	
	private static class InputArgs
	{
		public String eFile;
		public String fFile;
		public Integer modelNo;
		public Integer sentencesNo;
		public Integer iterations;
		public Boolean randomInit;
		public Boolean bothDirections;
		public String outputFile;
		public String tFile;
	}
}
